#李聪 2022/3/10
#生成哈希码
import os

import torch
import time
from PIL import Image
from torch.utils.data import DataLoader
from torchvision import transforms

from datasets import AID_Dataset
IMAGE_ROW = 4  # 图片间隔，也就是合并成一张图后，一共有几行
IMAGE_COLUMN = 5  # 图片间隔，也就是合并成一张图后，一共有几列
IMAGE_SAVE_PATH = 'final.jpg'  # 
IMAGE_RE_NUM=40 #呈现的图片数量

def hamming(a,b):
    distance=0
    for i in range(len(a)):
        if a[i]!=b[i]:
            distance=distance+1
    return distance
# 定义图像拼接函数
def image_compose(imglist):
    to_image = Image.new('RGB', (IMAGE_COLUMN * 224, IMAGE_ROW * 224))  # 创建一个新图
    # 循环遍历，把每张图片按顺序粘贴到对应位置上
    for y in range(1, IMAGE_ROW + 1):
        for x in range(1, IMAGE_COLUMN + 1):
            to_image.paste(imglist[5*(y-1)+x-1], ((x - 1) * 224, (y - 1) * 224))
    return to_image.save(IMAGE_SAVE_PATH)  # 保存新图
if __name__ == '__main__':
    device = torch.device("cuda")
    st_ti=time.time()
    print("正在加载系统....请稍等")
    print("预计需要一分半钟...")
    model_name = "sq_model_49"  # 参数代填
    model = torch.load(model_name)
    model.to(device)
    hash_list=[]
    root_dir = "../autodl-tmp/test"
    #root_dir = "../autodl-tmp/AID"
    #root_dir = "../AID"
    label_dir_list = os.listdir(root_dir)
    # print(label_dir_list)
    # input()
    AID = AID_Dataset(root_dir, label_dir_list)
    dataloader = DataLoader(dataset=AID, batch_size=64, shuffle=False)
    for data in dataloader:
        imgs,labels=data
        imgs=imgs.to(device)
        labels=labels.to(device)
        outputs=model(imgs)
        outputs_tem=torch.sign(outputs-0.5)
        hash_code=(outputs_tem+1)/2
        for i in range(len(imgs)):
            hash_list.append(hash_code[i])

    tran_totensor = transforms.ToTensor()
    tran_resize = transforms.Resize((224, 224))
    img_num = []
    for i in range(len(label_dir_list)):
        img_class_path = os.path.join(root_dir, label_dir_list[i])
        img_list = os.listdir(img_class_path)
        img_num.append(len(img_list))
    ed_ti=time.time()
    print("加载完成！")
    print("加载用时:{}".format(ed_ti-st_ti))
    find_num=0
    total_acc=0
    while True:
        img_name_find = input("输入图片名字:")
        img_label=input("输入图片类别:")
        img_path_qq=os.path.join(root_dir,img_label)
        img_path_end=os.path.join(img_path_qq,img_name_find)
        start_time=time.time()
        if os.path.exists(img_path_end)==False:
            print("{}不存在，请重新输入！".format(img_name_find))
            continue
        img=Image.open(img_path_end)
        img=img.convert('RGB')
        img=tran_resize(img)
        img=tran_totensor(img)
        img=img.view([1,3,224,224])
        img=img.to(device)
        out=model(img)
        re_list = []
        result=[]
        out = torch.sign(out - 0.5)
        hash = (out + 1) / 2
        #print(hash[0])
        #print(hash_list[0])
        for h in hash_list:
            tem=hamming(hash[0],h)
            re_list.append(tem)
        for i in range(IMAGE_RE_NUM):
            index=-1
            min=50
            for j in range(len(re_list)):
                if re_list[j]<min:
                    min=re_list[j]
                    index=j
            re_list[index]=50
            result.append(index)
        img_result_list=[]
        accuracy=0
        label_index_end=0
        print("检索结果:")
        for i in result:
            length = 0
            label_index = 0
            for s in range(len(label_dir_list)):
                length = length + img_num[s]
                if i < length:
                    label_index = s
                    break
            label_name=label_dir_list[label_index]
            img_class_path = os.path.join(root_dir, label_dir_list[label_index])
            img_list = os.listdir(img_class_path)
            img_name = img_list[i - length + img_num[label_index]]
            img_path = os.path.join(img_class_path, img_name)
            print(img_path)
            img_result=Image.open(img_path)
            img_result_list.append(img_result)
            if img_label==label_dir_list[label_index]:
                accuracy=accuracy+1
            label_index_end=label_index
        acc=accuracy/IMAGE_RE_NUM
        total_acc=total_acc+accuracy
        find_num=find_num+1
        print("本次检索准确率:{}".format(acc))
        for i in range(len(img_result_list)):
            img_result_list[i]=tran_resize(img_result_list[i])
        image_compose(img_result_list)
        HH=Image.open(IMAGE_SAVE_PATH)
        HH.show()
        end_time=time.time()
        print("检索用时:{}".format(end_time-start_time))
        print("总体的检索准确率:{}".format(total_acc/(find_num*IMAGE_RE_NUM)))